Add a "gradient zoo" page? #2447
Conversation
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## master #2447 +/- ##
===========================================
+ Coverage 33.20% 71.52% +38.31%
===========================================
Files 31 31
Lines 1843 1921 +78
===========================================
+ Hits 612 1374 +762
+ Misses 1231 547 -684 ☔ View full report in Codecov by Sentry. |
|
|
||
| * Long compilation times, on the first call. | ||
|
|
||
| * Does not at present work on all Flux models, due to missing rules. |
There was a problem hiding this comment.
Out of curiousity, what models does it not work on atm?
There was a problem hiding this comment.
FWIW all the CPU Flux tests I know of presently pass (and please let me know any failing ones so I can fix them), including some which Zygote fails on:
Flux.jl/test/ext_enzyme/enzyme.jl
Line 143 in d39d13b
There was a problem hiding this comment.
That's great!
Sorry this was written a few months back, and I don't recall what issues I ran into. I see many have been fixed.
|
|
||
| New package which works on the LLVM code which Julia compiles down to. | ||
|
|
||
| * Allows mutation of arrays. |
There was a problem hiding this comment.
I might also add high-performance :) [by virtue of pioneering the "run after optimization" trick (https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf)]
|
|
||
| * Does not always handle type instability. | ||
|
|
||
| * Custom rules by its own rules... Generally fewer such rules than Zygote, and at a lower level -- applied to `BLAS.gemm!` not `*`. |
There was a problem hiding this comment.
Enzyme also supports importing chainrules as well (though it is discouraged as they are usually lower-performance for sake of not understanding mutation)
| Zygote.hessian is like this. | ||
|
|
||
| ### Enzyme.jl | ||
|
|
There was a problem hiding this comment.
Haven't used on Flux functions, but it is well tested in our CI
A while back I tried to write up a docs page pointing out that Flux is AD-agnostic -- anything which produces an appropriately nested structure should work.
This is a rough draft, help with tidying it up would be welcome.
Perhaps it's too opinionated about other packages. Certainly any opinions it offers should have clear dates attached.
Rendered version: https://github.com/mcabbott/Flux.jl/blob/grad_zoo/docs/src/tutorials/gradient_zoo.md